#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <errno.h>
#include <limits.h>
#include <unistd.h>
#include <sys/types.h>
#include <fcntl.h>

#define CGROUP_ROOT "/sys/fs/cgroup"
#define CGROUP_NAME "my_cgroup"
#define MAX_PATH 1024

int mkdir_p(const char *path, mode_t mode) {
    char *ppath = strdup(path); // 复制路径（避免修改原字符串）
    char *sp = ppath; // 用于遍历路径的指针
    int status = 0;

    // 处理根目录（以 '/' 开头）
    if (*sp == '/') {
        sp++;
    }

    while (1) {
        sp = strpbrk(sp, "/\\"); // 支持 Windows 和 Linux 路径分隔符
        if (sp != NULL) {
            *sp = '\0'; // 临时截断路径
        }

        // 创建当前层级的目录
        if (mkdir(ppath, mode)){
            if (errno != EEXIST) { // 忽略目录已存在的错误
                perror("mkdir failed");
                status = -1;
                break;
            }
        }

        if (sp == NULL) {
            break; // 路径遍历完成
        }

        *sp = '/'; // 恢复路径分隔符
        sp++;
    }

    free(ppath);
    return status;
}


int create_multilevel_dir(const char *path) {
    char *path_copy = strdup(path);
    if (path_copy == NULL) {
        return -1; // 内存分配失败
    }

    char current_path[PATH_MAX];
    current_path[0] = '\0';

    size_t path_len = strlen(path_copy);
    char *p = path_copy;
    size_t current_pos = 0;

    // 处理绝对路径的情况
    if (path_len > 0 && path_copy[0] == '/') {
        strcpy(current_path, "/");
        current_pos = 1;
        p++; // 跳过开头的 '/'
    }

    while (*p != '\0') {
        // 跳过所有连续的 '/'
        while (*p == '/') {
            p++;
        }
        if (*p == '\0') {
            break; // 处理完毕
        }

        char *segment_start = p;
        while (*p != '/' && *p != '\0') {
            p++;
        }
        size_t segment_len = p - segment_start;

        // 检查路径长度是否超出缓冲区
        if (current_pos + segment_len + 2 > PATH_MAX) {
            free(path_copy);
            errno = ENAMETOOLONG;
            return -1;
        }

        // 添加目录分隔符（非起始位置）
        if (current_pos != 0) {
            current_path[current_pos++] = '/';
        }

        // 拼接当前目录段
        strncpy(current_path + current_pos, segment_start, segment_len);
        current_pos += segment_len;
        current_path[current_pos] = '\0';

        // 检查目录是否存在
        struct stat st;
        if (stat(current_path, &st) != 0) {
            // 目录不存在，创建
            if (mkdir(current_path, 0755) != 0) {
                free(path_copy);
                return -1;
            }
        } else if (!S_ISDIR(st.st_mode)) {
            // 存在但不是目录
            free(path_copy);
            errno = ENOTDIR;
            return -1;
        }
    }

    free(path_copy);
    return 0;
}


// 创建cgroup目录
static int create_cgroup_dir(const char *path) {
    if (mkdir(path, 0755) == -1) {
        if (errno != EEXIST) { // 忽略已存在的错误
            perror("mkdir failed");
            return -1;
        }
    }
    return 0;
}

// 将进程PID写入cgroup的cgroup.procs文件
static int write_pid_to_cgroup(const char *cgroup_path, pid_t pid) {
    char procs_path[MAX_PATH];
    snprintf(procs_path, MAX_PATH, "%s/cgroup.procs", cgroup_path);

    int fd = open(procs_path, O_WRONLY | O_APPEND);
    if (fd == -1) {
        perror("open cgroup.procs failed");
        return -1;
    }

    char pid_str[16];
    snprintf(pid_str, sizeof(pid_str), "%d", pid);

    if (write(fd, pid_str, strlen(pid_str)) == -1) {
        perror("write pid to cgroup.procs failed");
        close(fd);
        return -1;
    }

    close(fd);
    return 0;
}

// 设置内存限制（字节）
static int set_memory_limit(const char *cgroup_path, const char *limit) {
    char memory_max_path[MAX_PATH];
    snprintf(memory_max_path, MAX_PATH, "%s/memory.max", cgroup_path);

    int fd = open(memory_max_path, O_WRONLY);
    if (fd == -1) {
        perror("open memory.max failed");
        return -1;
    }

    if (write(fd, limit, strlen(limit)) == -1) {
        perror("write memory.max failed");
        close(fd);
        return -1;
    }

    close(fd);
    return 0;
}

// 设置CPU限制（格式: "max 100000"）
static int set_cpu_limit(const char *cgroup_path, const char *limit) {
    char cpu_max_path[MAX_PATH];
    snprintf(cpu_max_path, MAX_PATH, "%s/cpu.max", cgroup_path);

    int fd = open(cpu_max_path, O_WRONLY);
    if (fd == -1) {
        perror("open cpu.max failed");
        return -1;
    }

    if (write(fd, limit, strlen(limit)) == -1) {
        perror("write cpu.max failed");
        close(fd);
        return -1;
    }

    close(fd);
    return 0;
}


// int main(int argc, char *argv[]) {
//     if (argc != 3) {
//         fprintf(stderr, "Usage: %s <PID> <MEMORY_LIMIT_IN_MB>\n", argv[0]);
//         exit(EXIT_FAILURE);
//     }

//     pid_t pid = atoi(argv[1]);
//     int memory_limit_mb = atoi(argv[2]);
//     char memory_limit_str[32];
//     snprintf(memory_limit_str, sizeof(memory_limit_str), "%dM", memory_limit_mb);

//     // 构建cgroup路径
//     char cgroup_path[MAX_PATH];
//     snprintf(cgroup_path, MAX_PATH, "%s/%s", CGROUP_ROOT, CGROUP_NAME);

//     // 创建cgroup目录
//     if (create_cgroup_dir(cgroup_path) != 0) {
//         exit(EXIT_FAILURE);
//     }

//     // 将进程加入cgroup
//     if (write_pid_to_cgroup(cgroup_path, pid) != 0) {
//         exit(EXIT_FAILURE);
//     }

//     // 设置内存限制
//     if (set_memory_limit(cgroup_path, memory_limit_str) != 0) {
//         exit(EXIT_FAILURE);
//     }

//     // 设置CPU限制（示例：限制为1个CPU核心）
//     if (set_cpu_limit(cgroup_path, "100000 100000") != 0) {
//         exit(EXIT_FAILURE);
//     }

//     printf("Successfully added PID %d to cgroup %s\n", pid, cgroup_path);
//     printf("Memory limit set to %s\n", memory_limit_str);
//     printf("CPU limit set to 100000\n");

//     // 保持程序运行以便观察效果
//     printf("Press Enter to exit and clean up...\n");
//     getchar();

//     // 清理：删除cgroup（需要先移出所有进程）
//     // 注意：实际应用中需确保cgroup内无进程后再删除
//     if (rmdir(cgroup_path) == -1) {
//         perror("Failed to delete cgroup directory");
//     }

//     return 0;
// }